#!/usr/bin/env python3

import torch
from safetensors import safe_open
import os
import sys
import numpy as np

def analyze_adapter_model(file_path):
    """详细分析adapter_model.safetensors文件的内容"""
    print(f"详细分析文件: {file_path}")
    
    if not os.path.exists(file_path):
        print(f"文件不存在: {file_path}")
        return
    
    total_params = 0
    
    # 打开safetensors文件
    with safe_open(file_path, framework="pt", device="cpu") as f:
        keys = list(f.keys())
        
        print(f"\n总参数数量: {len(keys)}")
        
        # 按层分组并统计参数
        layer_stats = {}
        
        for key in keys:
            tensor = f.get_tensor(key)
            param_count = tensor.numel()
            total_params += param_count
            
            # 解析层信息
            parts = key.split('.')
            if len(parts) > 5:
                layer_num = parts[4]  # 层号
                if layer_num not in layer_stats:
                    layer_stats[layer_num] = {'count': 0, 'params': 0}
                layer_stats[layer_num]['count'] += 1
                layer_stats[layer_num]['params'] += param_count
        
        print(f"\n总参数量: {total_params:,} 个参数")
        print(f"\n每层参数分布:")
        for layer_num in sorted(layer_stats.keys(), key=lambda x: int(x) if x.isdigit() else x):
            stats = layer_stats[layer_num]
            print(f"  Layer {layer_num:>2}: {stats['count']:>2} 个参数矩阵, {stats['params']:>8,} 个参数")
        
        # 显示一些参数的统计信息
        print(f"\n参数值统计:")
        sample_keys = keys[:5]  # 只分析前5个参数矩阵
        for key in sample_keys:
            tensor = f.get_tensor(key)
            print(f"  {key}:")
            print(f"    形状: {list(tensor.shape)}")
            print(f"    最小值: {tensor.min().item():.6f}")
            print(f"    最大值: {tensor.max().item():.6f}")
            print(f"    平均值: {tensor.mean().item():.6f}")
            print(f"    标准差: {tensor.std().item():.6f}")

if __name__ == "__main__":
    if len(sys.argv) != 2:
        print("使用方法: python analyze_adapter.py <adapter_model.safetensors文件路径>")
        sys.exit(1)
        
    file_path = sys.argv[1]
    analyze_adapter_model(file_path)
